# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
import random
import torch
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig


def random_factor_pair(M):
    # 找出所有正因子
    factors = [i for i in range(1, M+1) if M % i == 0]
    # 随机选择一个因子A
    A = random.choice(factors)
    B = M // A
    return A, B

@GENERATOR_REGISTRY.register("generator_ascend_generate_moe_finalize_routing")
class AscendMoeFinalizeRoutingV2(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]]
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:

        return input_case

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        expandedX = case_config.inputs[0]
        # expandedRowIdx = case_config.inputs[1]
        x1 = case_config.inputs[1]
        x2Optional = case_config.inputs[2]
        bias = case_config.inputs[3]
        scales = case_config.inputs[4]
        expandedRowIdx = case_config.inputs[5]
        expandedExpertIdx = case_config.inputs[6]
        out = case_config.inputs[7]
        # print(expandedRowIdx.range_values)
        E = bias.shape[0]
        Num_Rows = out.shape[0]
        H = expandedX.shape[1]
        K = scales.shape[1]
        if K > E:
            K = E
            scales.shape[1] = K # K: 为从总的专家E中选出K个专家;  E: expert num，即专家数，E需要大于等于K; 
        expandedX.shape = [Num_Rows * K, H] # drop less 场景shape为（NUM\_ROWS \* K, H）
        expandedRowIdx.shape = [Num_Rows * K]
        expandedRowIdx.range_values = [0, Num_Rows * K - 1] # dropPadMode参数值为0、2时，Tensor中的值取值范围是[0,NUM\_ROWS \* K-1]; 
        x1.shape = [Num_Rows, H] # x1 shape要求与out的shape一致
        x1.dtype = expandedX.dtype #x1 数据类型要求与expandedX一致
        x2Optional.shape = [Num_Rows, H] # x2Optional shape要求与out的shape一致
        x2Optional.dtype = expandedX.dtype # x2Optional 数据类型要求与expandedX一致
        bias.shape = [E, H] # 限制：其shape支持（E，H）
        bias.dtype = expandedX.dtype # 数据类型要求与expandedX一致
        scales.shape = [Num_Rows, K] # 限制：其shape支持（NUM\_ROWS，K）
        scales.dtype = expandedX.dtype # 数据类型要求与expandedX一致
        expandedExpertIdx.shape = [Num_Rows, K] # 限制：其shape支持（NUM\_ROWS，K）
        expandedExpertIdx.range_values = [0, E -1] # 限制：Tensor中的值取值范围是[0, E-1]
        out.dtype = expandedX.dtype
        # print(case_config)
        return case_config